package com.example.devsa_platform.service.impl;

import co.elastic.clients.elasticsearch._types.aggregations.Aggregation;
import co.elastic.clients.elasticsearch._types.aggregations.Buckets;
import co.elastic.clients.elasticsearch._types.aggregations.StringTermsBucket;
import com.example.devsa_platform.common.Result;
import com.example.devsa_platform.pojo.doc.BaseInfoDoc;
import com.example.devsa_platform.pojo.vo.stats.BaseInfoStatsVo;
import com.example.devsa_platform.service.BaseInfoService;
import jakarta.annotation.Resource;
import org.springframework.data.elasticsearch.client.elc.ElasticsearchAggregations;
import org.springframework.data.elasticsearch.client.elc.ElasticsearchTemplate;
import org.springframework.data.elasticsearch.client.elc.NativeQuery;
import org.springframework.data.elasticsearch.core.AggregationsContainer;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.stream.Collectors;

@Service
public class BaseInfoServiceImpl implements BaseInfoService {

    @Resource
    private ElasticsearchTemplate template;

    @Override
    public Result<BaseInfoStatsVo> baseInfoStats() {
        NativeQuery query = NativeQuery.builder()
                .withAggregation("main_branch_group", Aggregation.of(ag -> ag
                        .terms(te -> te.field("main_branch.keyword"))))
                .withAggregation("age_group", Aggregation.of(ag -> ag
                        .terms(te -> te.field("age.keyword"))))
                .withAggregation("employment_group", Aggregation.of(ag -> ag
                        .terms(te -> te.field("employment.keyword"))))
                .withAggregation("remote_work_group", Aggregation.of(ag -> ag
                        .terms(te -> te.field("remote_work.keyword"))))
                .withAggregation("coding_activities_group", Aggregation.of(ag -> ag
                        .terms(te -> te.field("coding_activities.keyword"))))
                .build();
        AggregationsContainer<?> aggregationsContainer = template.search(query, BaseInfoDoc.class).getAggregations();
        BaseInfoStatsVo baseInfoStatsVo = null;
        if (aggregationsContainer instanceof ElasticsearchAggregations elasticsearchAggregations) {
            List<Map<String,Object>> mainBranchGroup = parseAndAggregateResults("main_branch_group",elasticsearchAggregations);
            List<Map<String,Object>> ageGroup = parseAndAggregateResults("age_group",elasticsearchAggregations);
            List<Map<String,Object>> employmentGroup = parseAndAggregateResults("employment_group",elasticsearchAggregations);
            List<Map<String,Object>> remoteWorkGroup = parseAndAggregateResults("remote_work_group",elasticsearchAggregations);
            List<Map<String,Object>> codingActivitiesGroup = parseAndAggregateResults("coding_activities_group",elasticsearchAggregations);
            baseInfoStatsVo = new BaseInfoStatsVo(mainBranchGroup,ageGroup,employmentGroup,remoteWorkGroup,codingActivitiesGroup);
        }
        return Result.success(baseInfoStatsVo);
    }

    private List<Map<String, Object>> parseAndAggregateResults(String aggregationName, ElasticsearchAggregations elasticsearchAggregations) {
        Map<String, Long> aggregatedCounts = new HashMap<>();

        List<Map<String, Object>> originalResults = parseAggregation(aggregationName, elasticsearchAggregations);
        for (Map<String, Object> result : originalResults) {
            String key = (String) result.get("key");
            Long docCount = ((Number) result.get("doc_count")).longValue();

            // Split the key by ";"
            String[] parts = key.split(";");
            for (String part : parts) {
                part = part.trim();
                aggregatedCounts.put(part, aggregatedCounts.getOrDefault(part, 0L) + docCount);
            }
        }

        // Convert aggregatedCounts to the desired output
        List<Map<String, Object>> parsedResults = new ArrayList<>();
        for (Map.Entry<String, Long> entry : aggregatedCounts.entrySet()) {
            Map<String, Object> map = new HashMap<>();
            map.put("key", entry.getKey());
            map.put("doc_count", entry.getValue());
            parsedResults.add(map);
        }

        return parsedResults;
    }

    private List<Map<String, Object>> parseAggregation(String aggName, ElasticsearchAggregations elasticsearchAggregations) {
        Buckets<StringTermsBucket> buckets = Objects.requireNonNull(elasticsearchAggregations.get(aggName)).aggregation().getAggregate().sterms().buckets();
        return buckets.array().stream()
                .map(bucket -> {
                    String key = bucket.key().stringValue();
                    long docCount = bucket.docCount();
                    Map<String, Object> map = new HashMap<>();
                    map.put("key", key);
                    map.put("doc_count", docCount);
                    return map;
                })
                .collect(Collectors.toList());
    }
}
